import math
import random

def euclidean_distance(point1, point2):
    #计算欧几里得距离
    return math.sqrt(sum((a - b) ** 2 for a, b in zip(point1, point2)))

def assign_cluster(X, centers):
    # X: list of lists, 数据点列表 centers: list of lists, 聚类中心列表 labels: 每个样本所属的簇索引
    labels = []
    for point in X:
        min_dist = float('inf')
        best_center = 0
        for j, center in enumerate(centers):
            dist = euclidean_distance(point, center)
            if dist < min_dist:
                min_dist = dist
                best_center = j
        labels.append(best_center)
    return labels

def kmeans(X, k, epsilon=1e-4, max_iterations=100):
#X: 数据集 k: 聚类数 epsilon: 收敛阈值 max_iterations: 最大迭代次数 centers: 聚类中心 labels: 聚类标签

    # 随机初始化中心
    centers = random.sample(X, k)
    
    for iteration in range(max_iterations):
        # 分配簇
        labels = assign_cluster(X, centers)
        
        # 更新中心
        new_centers = []
        for j in range(k):
            # 获取属于该簇的所有点
            cluster_points = [X[i] for i in range(len(X)) if labels[i] == j]
            
            if len(cluster_points) > 0:
                # 计算新中心（均值）
                n_features = len(X[0])
                new_center = [
                    sum(point[d] for point in cluster_points) / len(cluster_points)
                    for d in range(n_features)
                ]
                new_centers.append(new_center)
            else:
                # 空簇处理：重新随机初始化
                new_centers.append(random.choice(X))
        
        # 检查收敛
        center_shift = sum(
            euclidean_distance(centers[j], new_centers[j]) 
            for j in range(k)
        )
        
        if center_shift < epsilon:
            break
            
        centers = new_centers
    
    return centers, labels